from flask import Flask, g, jsonify
from flask_auth_mgr import AuthMgr

#--------------------#
app = Flask(__name__)
app.config['SECRET_KEY'] = 'the quick brown fox jumps over the lazy dog'
auth = AuthMgr(app.config['SECRET_KEY'])
#--------------------#

#--------------------#
users = {
    "john": auth.decode_password("hello"),
    "susan": auth.decode_password("bye")
}

tokens = {
    'refresh_token': {},
    'access_token': {}
}
#--------------------#

#--------------------#
@auth.set_save_token
def save_token(tdata):
    user = tdata['data']['user']
    tp = tdata['type']
    tokens[tp][user] = tdata['token']
    return True
@auth.set_get_token
def get_token(tdata):
    user = tdata['data']['user']
    tp = tdata['type']
    return tokens[tp][user]
#--------------------#

#--------------------#
@auth.verify_password
def verify_password(user, pwd):
    g.user = None
    if user not in users:
        return False
    if auth.check_password(users.get(user), pwd):
        g.user = user
        return True
    return False

@auth.verify_refresh_token
def verify_rftoken(token):
    data = auth.refresh_token_data(token)
    if data is None:
        return False
    g.user = data['user']
    return True

@auth.verify_access_token
def verify_atoken(token):
    data = auth.access_token_data(token)
    if data is None:
        return False
    g.user = data['user']
    return True
#--------------------#

#--------------------#
@app.route('/refresh_token')
@auth.required_login
def route_refresh_token():
    token_data = {
        'user': g.user,
    }

    data = {
        'refresh_token': auth.generate_refresh_token(token_data),
    }
    return jsonify(data)

@app.route('/access_token')
@auth.required_refresh_token
def route_access_token():
    token_data = {
        'user': g.user,
    }

    data = {
        'access_token': auth.generate_access_token(token_data),
    }
    return jsonify(data)

# rftoken -> rftoken
@app.route('/token')
@auth.required_multi_refresh
def route_token():
    token_data = {
        'user': g.user,
    }

    data = {
        'refresh_token': auth.generate_refresh_token(token_data),
        'access_token': auth.generate_access_token(token_data),
    }
    return jsonify(data)
#--------------------#
@app.route('/need_pwd')
@auth.required_login
def route_need_pwd():
    return 'this page need password: %s\n' % g.user

@app.route('/need_rftoken')
@auth.required_refresh_token
def route_need_rftoken():
    return 'this page need refresh token: %s\n' % g.user

@app.route('/need_pwd_rftoken')
@auth.required_multi_refresh
def route_need_pwd_rftoken():
    return 'this page need password or refresh token: %s\n' % g.user

@app.route('/need_atoken')
@auth.required_access_token
def route_need_atoken():
    return 'this page need access token: %s\n' % g.user

@app.route('/need_pwd_atoken')
@auth.required_multi_access
def route_need_pwd_atoken():
    return 'this page need password or access token: %s\n' % g.user
#--------------------#

#--------------------#
if __name__ == '__main__':
    app.run(debug=True)
#--------------------#
